#include "../inc/ff_utils.h"
#include "../inc/ff_utils_cuda.cuh"
#include <assert.h>

using namespace feifei;


typedef enum class MemTypeEnum
{
	Page = 1,
	Pin = 2,
	WrCmb = 3,
	Map = 4,
	Pitch = 5,
	Dev = 6
}E_MemType;
typedef enum class BwDir
{
	H2D = 1,
	D2H = 2,
	D2D = 3,
	H2H = 4
}E_BwDir;
void bandwidthTest(size_t width, size_t height, E_MemType mem, bool isPitch, BwDir dir, bool isAsync)
{
	size_t pitch = width * sizeof(float);
	size_t dataSize = width * height;
	size_t memSize = dataSize * sizeof(float);
	int iteration = 100;
	float elapsedTimeInMs = 0.0f;
	float bandwidthInMBs = 0.0f;

	float* h_data1, * h_data2;
	float* d_data1, * d_data2;

	cudaError_t err;
	cudaStream_t stream;
	cudaEvent_t startEvt, stopEvt;
	err = cudaStreamCreate(&stream); assert(err == cudaSuccess);
	err = cudaEventCreate(&startEvt, 0); assert(err == cudaSuccess);
	err = cudaEventCreate(&stopEvt, 0); assert(err == cudaSuccess);

	// 1. allocate device memory
	if (isPitch)
	{
		err = cudaMallocPitch(&d_data1, &pitch, width * sizeof(float), height); assert(err == cudaSuccess);
		err = cudaMallocPitch(&d_data2, &pitch, width * sizeof(float), height); assert(err == cudaSuccess);
		memSize = pitch * height;
	}
	else
	{
		err = cudaMalloc(&d_data1, memSize); assert(err == cudaSuccess);
		err = cudaMalloc(&d_data2, memSize); assert(err == cudaSuccess);
	}

	// 2. allocate host memory
	switch (mem)
	{
	case E_MemType::Page:
		h_data1 = (float*)malloc(memSize); assert(h_data1 != nullptr);
		h_data2 = (float*)malloc(memSize); assert(h_data2 != nullptr);
		break;
	case E_MemType::Pin:
		err = cudaMallocHost((void**)&h_data1, memSize, cudaHostAllocDefault); assert(err == cudaSuccess);
		err = cudaMallocHost((void**)&h_data2, memSize, cudaHostAllocDefault); assert(err == cudaSuccess);
		break;
	case E_MemType::WrCmb:
		err = cudaMallocHost((void**)&h_data1, memSize, cudaHostAllocWriteCombined); assert(err == cudaSuccess);
		err = cudaMallocHost((void**)&h_data2, memSize, cudaHostAllocWriteCombined); assert(err == cudaSuccess);
		break;
	case E_MemType::Map:
		err = cudaMallocHost((void**)&h_data1, memSize, cudaHostAllocMapped); assert(err == cudaSuccess);
		err = cudaMallocHost((void**)&h_data2, memSize, cudaHostAllocMapped); assert(err == cudaSuccess);
		break;
	default:
		h_data1 = (float*)malloc(memSize); assert(h_data1 != nullptr);
		h_data2 = (float*)malloc(memSize); assert(h_data2 != nullptr);
	}

	// 3. initialize memory
	for (uint64_t i = 0; i < memSize / sizeof(float); i++)
		h_data1[i] = (i & 0xff) * 1.0f;
	if (dir == BwDir::D2H || dir == BwDir::D2D)
	{
		err = cudaMemcpy(d_data1, h_data1, memSize, cudaMemcpyHostToDevice);
		assert(err == cudaSuccess);
	}

	// 4. copy data
	err = cudaEventRecord(startEvt, stream);
	for (int i = 0; i < iteration; i++)
	{
		if (isPitch)
		{
			if (isAsync)
			{
				switch (dir)
				{
				case BwDir::H2D:	err = cudaMemcpy2DAsync(d_data1, pitch, h_data1, pitch, width, height, cudaMemcpyHostToDevice); assert(err == cudaSuccess); break;
				case BwDir::D2H:	err = cudaMemcpy2DAsync(d_data1, pitch, h_data1, pitch, width, height, cudaMemcpyDeviceToHost); assert(err == cudaSuccess); break;
				case BwDir::D2D:	err = cudaMemcpy2DAsync(d_data1, pitch, h_data1, pitch, width, height, cudaMemcpyDeviceToDevice); assert(err == cudaSuccess); break;
				}
			}
			else
			{
				switch (dir)
				{
				case BwDir::H2D:	err = cudaMemcpy2D(d_data1, pitch, h_data1, pitch, width, height, cudaMemcpyHostToDevice); assert(err == cudaSuccess); break;
				case BwDir::D2H:	err = cudaMemcpy2D(d_data1, pitch, h_data1, pitch, width, height, cudaMemcpyDeviceToHost); assert(err == cudaSuccess); break;
				case BwDir::D2D:	err = cudaMemcpy2D(d_data1, pitch, h_data1, pitch, width, height, cudaMemcpyDeviceToDevice); assert(err == cudaSuccess); break;
				}
			}
		}
		else
		{
			if (isAsync)
			{
				switch (dir)
				{
				case BwDir::H2D:	err = cudaMemcpyAsync(d_data1, h_data1, memSize, cudaMemcpyHostToDevice, stream); assert(err == cudaSuccess); break;
				case BwDir::D2H:	err = cudaMemcpyAsync(h_data2, d_data1, memSize, cudaMemcpyDeviceToHost, stream); assert(err == cudaSuccess); break;
				case BwDir::D2D:	err = cudaMemcpyAsync(d_data2, d_data1, memSize, cudaMemcpyDeviceToDevice, stream); assert(err == cudaSuccess); break;
				}
			}
			else
			{
				switch (dir)
				{
				case BwDir::H2D:	err = cudaMemcpy(d_data1, h_data1, memSize, cudaMemcpyHostToDevice); assert(err == cudaSuccess); break;
				case BwDir::D2H:	err = cudaMemcpy(h_data2, d_data1, memSize, cudaMemcpyDeviceToHost); assert(err == cudaSuccess); break;
				case BwDir::D2D:	err = cudaMemcpy(d_data2, d_data1, memSize, cudaMemcpyDeviceToDevice); assert(err == cudaSuccess); break;
				}
			}
		}
	}

	//get the total elapsed time in ms
	err = cudaEventRecord(stopEvt, stream);
	err = cudaEventSynchronize(stopEvt);
	err = cudaEventElapsedTime(&elapsedTimeInMs, startEvt, stopEvt);

	//calculate bandwidth in MB/s
	bandwidthInMBs = ((float)(1 << 10) * memSize * (float)iteration) / (elapsedTimeInMs * (float)(1 << 20));

	if (mem == E_MemType::Page)
	{
		free(h_data1);
		free(h_data2);
	}
	else
	{
		cudaFreeHost(h_data1);
		cudaFreeHost(h_data2);
	}
	cudaFree(d_data1);
	cudaFree(d_data2);
	cudaEventDestroy(startEvt);
	cudaEventDestroy(stopEvt);
	cudaStreamDestroy(stream);

	std::string logstr;
	if (isAsync)logstr = "Async";
	else		logstr = "Sync ";
	switch (dir)
	{
	case BwDir::H2D:	logstr += " H -> D"; break;
	case BwDir::D2H:	logstr += " D -> H"; break;
	case BwDir::D2D:	logstr += " D -> D"; break;
	}
	if (isPitch)
	{
		logstr += " Pitch  ";
	}
	else
	{
		logstr += " UnPitch";
	}
	if (dir == BwDir::D2D)
	{
		logstr += "         ";
	}
	else
	{
		switch (mem)
		{
		case E_MemType::Page:	logstr += " pageable"; break;
		case E_MemType::Pin:	logstr += " Pinned  "; break;
		case E_MemType::WrCmb:	logstr += " wr cmb  "; break;
		case E_MemType::Map:	logstr += " mapped  "; break;
		}
	}
	LOG("%s %d(MB) = %.3f(GB/s).", logstr.data(), memSize / 1024 / 1024, bandwidthInMBs / 1024);
}
void BandwidthTest()
{
	size_t width = 16 * 1024;
	size_t height = 2 * 1024;
	PrintSeperator('=');
	INFO("BandWidth Test");

	PrintSeperator('-');
	LOG("H->D VS D->H VS D->D");
	bandwidthTest(width, height, E_MemType::Pin, false, BwDir::H2D, false);
	bandwidthTest(width, height, E_MemType::Pin, false, BwDir::D2H, false);
	bandwidthTest(width, height, E_MemType::Pin, false, BwDir::D2D, false);
	PrintSeperator('-');
	LOG("Pageable VS Pinned");
	bandwidthTest(width, height, E_MemType::Page, false, BwDir::H2D, false);
	bandwidthTest(width, height, E_MemType::Pin, false, BwDir::H2D, false);
	PrintSeperator('-');
	LOG("Write Combined");
	bandwidthTest(width, height, E_MemType::Pin, false, BwDir::H2D, false);
	bandwidthTest(width, height, E_MemType::WrCmb, false, BwDir::H2D, false);
	PrintSeperator('-');
	LOG("Sync VS Async");
	bandwidthTest(width, height, E_MemType::Pin, false, BwDir::H2D, true);
	bandwidthTest(width, height, E_MemType::Pin, false, BwDir::H2D, false);
	PrintSeperator('-');
	width = 16 * 1024 + 1023; height = 2 * 1024;
	LOG("Pintch Memory: width = %d(K), height = %d(K).", width / 1024, height / 1024);
	bandwidthTest(width, height, E_MemType::Pin, true, BwDir::H2D, false);
	bandwidthTest(width, height, E_MemType::Pin, false, BwDir::H2D, false);

	PrintSeperator('=');
}

int TestMemoryDump()
{
	PrintSeperator('-');
	PrintHeader("device memroy data test");

	float* h_a, * h_b;
	float* d_a, * d_b;

	size_t len = 128;
	size_t sz = len * sizeof(float);

	h_a = (float*)malloc(sz);
	h_b = (float*)malloc(sz);
	cudaMalloc(&d_a, sz);
	cudaMalloc(&d_b, sz);

	for (uint32_t i = 0; i < len; i++)
	{
		h_a[i] = 10.0f * i;
		h_b[i] = 0;
	}
	cudaMemcpy(d_a, h_a, sz, cudaMemcpyHostToDevice);

	DumpMemDev(d_a, len, get_debug_path() + "d_a_bin.dat");
	LoadMemDev((void**)(&d_b), get_debug_path() + "d_a_bin.dat");

	cudaMemcpy(h_b, d_b, sz, cudaMemcpyDeviceToHost);
	LogDataDev<float>(d_b, len, "d_b");

	E_ffState state = E_ffState::SUCCESS;
	for (uint32_t i = 0; i < len; i++)
	{
		if (h_b[i] != h_a[i])
		{
			state = E_ffState::RTN_ERR;
			LOG("compare failed:\n");
			LOG("[%d]: host = %.1f, device = %.1f\n", i, h_b[i], h_a[i]);
		}
	}

	free(h_a);
	free(h_b);
	cudaFree(d_a);
	cudaFree(d_b);

	return state == E_ffState::SUCCESS ? 0 : 1;
}

int main(int argc, char* argv[])
{
	int state = 0;

	PrintSeperator('*');
	PrintHeader("cuda utils test", '*');
	PrintSeperator('*');
	
	InitGpuRuntime();
	BandwidthTest();
	state += TestMemoryDump();

	LOG("error_cnt = %d", state);
	return state;
}
